package main

import (
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/rlp"
)

type ExtensionNode struct {
	Path []Nibble
	Next Node
}

func NewExtensionNode(nibbles []Nibble, next Node) *ExtensionNode {
	return &ExtensionNode{
		Path: nibbles,
		Next: next,
	}
}

func (e ExtensionNode) Hash() []byte {
	return crypto.Keccak256(e.Serialize())
}

func (e ExtensionNode) Raw() []interface{} {
	return e.RawWithDepth(0)
}

func (e ExtensionNode) RawWithDepth(depth int) []interface{} {
	hashes := make([]interface{}, 2)
	hashes[0] = ToBytes(ToPrefixed(e.Path, false))
	
	// 防止过深的递归调用
	if depth > 100 {
		// 如果递归深度过深，直接使用哈希值
		hashes[1] = e.Next.Hash()
		return hashes
	}
	
	// 根据节点类型决定如何处理
	switch next := e.Next.(type) {
	case *ExtensionNode:
		nextRaw := next.RawWithDepth(depth + 1)
		nextSerialized, err := rlp.EncodeToBytes(nextRaw)
		if err != nil || len(nextSerialized) >= 32 {
			hashes[1] = e.Next.Hash()
		} else {
			hashes[1] = nextRaw
		}
	case *BranchNode:
		nextRaw := next.RawWithDepth(depth + 1)
		nextSerialized, err := rlp.EncodeToBytes(nextRaw)
		if err != nil || len(nextSerialized) >= 32 {
			hashes[1] = e.Next.Hash()
		} else {
			hashes[1] = nextRaw
		}
	default:
		// 对于 LeafNode 或其他节点类型，直接调用 Raw()
		nextRaw := e.Next.Raw()
		nextSerialized, err := rlp.EncodeToBytes(nextRaw)
		if err != nil || len(nextSerialized) >= 32 {
			hashes[1] = e.Next.Hash()
		} else {
			hashes[1] = nextRaw
		}
	}
	
	return hashes
}

func (e ExtensionNode) Serialize() []byte {
	return Serialize(e)
}
